from ctypes import *
import cv2
import numpy as np
import numpy.ctypeslib as npct
import time
import math
from PIL import Image
print("cv2 version is:",cv2.__version__)
# -----------------------------------双目相机的基本参数---------------------------------------------------------
#   left_camera_matrix          左相机的内参矩阵c:\ProgramData\Microsoft\Windows\Start Menu\Programs\FileZilla FTP Client\FileZilla.lnk
#   right_camera_matrix         右相机的内参矩阵
#
#   left_distortion             左相机的畸变系数    格式(K1,K2,P1,P2,0)
#   right_distortion            右相机的畸变系数
# -------------------------------------------------------------------------------------------------------------
#内参

left_camera_matrix = np.array([[335.73514791566,
                                -0.731874511921937,
                                322.909744272018],
                                [0,
                                335.55542827028,
                                243.380676402428],
                                [0,0,1]])
right_camera_matrix = np.array([[336.040359146416,
                                -0.314558518222512,
                                345.151908479252],
                                [0,
                                335.596735734494,
                                235.933100504815],
                                [0,0,1]])
#畸变系数
left_distortion = np.array([[-0.0671729884111311,
                            -0.00226867715256577,
                            0.00159660436315177,
                            0.00229071232872235,
                            0]])
right_distortion = np.array([[-0.0681254193650723,
                            0.00284849515254842,
                            0.00137186177789993,
                            0.00172620053631537,
                            0]])
#旋转矩阵
R = np.array([[0.999365346521295,
                0.000290377473121344,
                -0.0356204976564704],
                [-0.0000372010398693062,
                0.999974736715963,
                0.0071080620370507],
                [0.0356216617868137,
                -0.00710222576119855,
                0.999340110073034]])
#平移矩阵
T = np.array([[-60.5633077641667], 
                  [0.0346502867236554], 
                  [-1.33159248369833]])
size = (640, 480)

R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(left_camera_matrix, left_distortion,
                                                                  right_camera_matrix, right_distortion, size, R,
                                                                  T)
# 校正查找映射表,将原始图像和校正后的图像上的点一一对应起来
left_map1, left_map2 = cv2.initUndistortRectifyMap(left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
right_map1, right_map2 = cv2.initUndistortRectifyMap(right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)
print(Q)
Q = np.array([
    [1.00000000e+00, 0.00000000e+00, 0.00000000e+00, -6.99876471e+02],
    [0.00000000e+00, 1.00000000e+00, 0.00000000e+00, -3.64293882e+02],
    [0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 6.67635587e+02],
    [0.00000000e+00, 0.00000000e+00, 1.64479769e-02, 0.00000000e+00]
    ])
# ---------------------------------------------------------------------------------------------------------
#   classes        coco数据集的种类,网络返回‘0’时，对应着person，依次类推
# ---------------------------------------------------------------------------------------------------------
classes = ('person','bicycle','car','motorbike','aeroplane','bus','train','truck','boat','traffic light',
'fire hydrant','stop sign','parking meter','bench','bird','cat','dog','horse','sheep','cow','elephant',
'bear','zebra','giraffe','backpack','umbrella','handbag','tie','suitcase','frisbee','skis','snowboard',
'sports ball','kite','baseball bat','baseball glove','skateboard','surfboard','tennis racket','bottle',
'wine glass','cup','fork','knife','spoon','bowl','banana','apple','sandwich','orange','broccoli','carrot',
'hot dog','pizza','donut','cake','chair','sofa','pottedplant','bed','diningtable','toilet','tvmonitor',
'laptop','mouse','remote','keyboard','cell phone','microwave','oven','toaster','sink','refrigerator',
'book','clock','vase','scissors','teddy bear','hair drier','toothbrush')


# ---------------------------------------------------------------------------------------------------------
#   Detector()       配置tensorrt加速
# ---------------------------------------------------------------------------------------------------------
class Detector():
    def __init__(self,model_path,dll_path):
        self.yolov5 = CDLL(dll_path)
        self.yolov5.Detect.argtypes = [c_void_p,c_int,c_int,POINTER(c_ubyte),npct.ndpointer(dtype = np.float32, ndim = 2, shape = (50, 6), flags="C_CONTIGUOUS")]
        self.yolov5.Init.restype = c_void_p
        self.yolov5.Init.argtypes = [c_void_p]
        self.yolov5.cuda_free.argtypes = [c_void_p]
        self.c_point = self.yolov5.Init(model_path)

    def predict(self,img):
        rows, cols = img.shape[0], img.shape[1]
        res_arr = np.zeros((50,6),dtype=np.float32)
        self.yolov5.Detect(self.c_point,c_int(rows), c_int(cols), img.ctypes.data_as(POINTER(c_ubyte)),res_arr)
        self.bbox_array = res_arr[~(res_arr==0).all(1)]
        return self.bbox_array

    def free(self):
        self.yolov5.cuda_free(self.c_point)

# ------------------------------------visualize可视化程序-------------------------------------------------------
#   img                     输入的图片
#   bbox_array              多组yolo网络预测的结果
#   middle_x、middle_y       检测目标的中心点坐标，用于测出距离distance
# -------------------------------------------------------------------------------------------------------------
def visualize(img,bbox_array):
    people_res = []
    for person in bbox_array:
        if int(person[4]) == 0:
            people_res.append(person)
    for temp in people_res:
        bbox = [temp[0],temp[1],temp[2],temp[3]]  #xywh
        clas = int(temp[4])
        score = temp[5]
        middle_x = int(np.floor(temp[0]+temp[2]*0.5))
        middle_y = int(np.floor(temp[1]+temp[3]*0.5))
        distance = math.sqrt(threeD[middle_y][middle_x][0] ** 2 +
                             threeD[middle_y][middle_x][1] ** 2 + threeD[middle_y][middle_x][2] ** 2)
        distance = distance / 1000.0  # mm -> m
        cv2.rectangle(img,(int(temp[0]),int(temp[1])),(int(temp[0]+temp[2]),int(temp[1]+temp[3])), (0, 0, 225), 2)
        img = cv2.putText(img, str(classes[int(clas)])+" "+str(round(score,2))+" dis="+str(round(distance,2)),
                          (int(temp[0]),int(temp[1])-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 225), 2)
    return img
def puttext(img, text):
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1
    font_color = (0, 255, 0)  
    thickness = 4
    text_size = cv2.getTextSize(text, font, font_scale, thickness)[0]
    text_width, text_height = text_size
    x = 10
    y = 60
    cv2.putText(img, text, (x, y), font, font_scale, font_color, thickness)
    return img
#---------------------------------------------------#
#   对输入图像进行不失真resize
#---------------------------------------------------#
def resize_image(image, size):
    iw, ih  = image.size
    w, h    = size
    scale   = min(w/iw, h/ih)
    nw      = int(iw*scale)
    nh      = int(ih*scale)

    image   = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128,128,128))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image

#----------------------加载推理所需要的engine、dll文件---------------------------------#
#   yolov5s_480         图片推理的输入格式为640x480
#   yolov5              图片推理的输入格式为640x640，若使用它，需要打开182行的resize函数
#------------------------------------------------------------------------------------#
if __name__ == "__main__":
    det = Detector(model_path=b"./yolov5n.engine",dll_path="./libyolov5.so")  # b'' is needed
    # 加载视频文件
    print("success load engine")
    cap = cv2.VideoCapture(-1)
    print("success open video")
    WIN_NAME = 'Deep disp'
    cv2.namedWindow(WIN_NAME, cv2.WINDOW_AUTOSIZE)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    width = 640
    height = 480
    print("width:%d,height:%d"%(width,height))
    
    LeftTop = (int(width / 2) - 50,int(height / 2) - 50)
    RightBottom = (int(width / 2) + 50,int(height / 2) + 50)  
    path = 'hsv_data.txt'
    wait_begin_time = time.time()
    #等待五秒
    while True:
        ret, img = cap.read()
        if ret:
            frame2 = img[0:, 0:640]
            frame1 = img[0:, 640:1280]
            wait_end_time = time.time()
            img = puttext(frame2,"waiting last time:%.2f"%(5 - (wait_end_time - wait_begin_time)))
            cv2.rectangle(frame2, LeftTop, RightBottom, (0,255,0),2, 4)
            if wait_end_time - wait_begin_time > 5:
                break
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            #img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
            cv2.imshow("frame2",frame2)
    cv2.destroyAllWindows()
    #学习颜色阈值
    learn_begin_time = time.time()
    H = []
    S = []
    V = []
    for i in range(100):
        t1 = time.time()
        ret,img = cap.read()
        if ret:
            frame2 = img[0:, 0:640]
            frame1 = img[0:, 640:1280]
            frame2 = cv2.cvtColor(frame2,cv2.COLOR_BGR2HSV)
            roi = frame2[LeftTop[1]:RightBottom[1],LeftTop[0]:RightBottom[0]]
            h, s, v = cv2.split(roi)
            H.append(h)
            S.append(s)
            V.append(v)
            time.sleep(0.00000001)
            t2 = time.time()
            fps = 1 / (t2 - t1)
            frame2 = cv2.cvtColor(frame2,cv2.COLOR_HSV2BGR)
            frame2 = puttext(frame2,"learning color -- FPS:%.2f"%fps)
            cv2.rectangle(frame2, LeftTop, RightBottom, (0,255,0),2, 4)
            #frame2 = cv2.cvtColor(frame2,cv2.COLOR_BGR2GRAY)
            cv2.imshow("frame2",frame2)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    cv2.destroyAllWindows()
    H_data = np.ravel(H)
    S_data = np.ravel(S)
    V_data = np.ravel(V)
    learn_end_time = time.time()
    print("learning use time:%.2f,one circle use time:%.5f"%((learn_end_time - learn_begin_time),(learn_end_time - learn_begin_time)/100))

    low_threshold = 5
    high_threshold = 95
    learn_start_time = time.time()
    H_low = np.percentile(H_data, low_threshold)
    H_high = np.percentile(H_data, high_threshold)
    S_low = np.percentile(S_data, low_threshold)
    S_high = np.percentile(S_data, high_threshold)
    V_low = np.percentile(V_data, low_threshold)
    V_high = np.percentile(V_data, high_threshold)
    
    H_threshold = [H_low, H_high]
    S_threshold = [S_low, S_high]
    V_threshold = [V_low, V_high]
    
    print(H_threshold)
    print(S_threshold)
    print(V_threshold)

    print("begin detect")
    #阈值上下界
    lower = np.array([H_threshold[0], S_threshold[0], V_threshold[0]])
    upper = np.array([H_threshold[1], S_threshold[1], V_threshold[1]])

    # 3 读取视频
    fps = 0.0
    ret, frame = cap.read()
    if ret:
        while True:
            # 是否读取到了帧，读取到了则为True
            ret, frame = cap.read()
            # 开始计时，用于计算帧率
            t1 = time.time()
            img_color = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame1 = frame[0:480, 0:640]
            frame2 = frame[0:480, 640:1280]  # 割开双目图像
            #-------------------------------------颜色阈值----------------------------------------------------------
            hsv_frame1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2HSV)
            hsv_frame2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2HSV)
            #掩膜
            mask1 = cv2.inRange(hsv_frame1, lower, upper)
            mask2 = cv2.inRange(hsv_frame2, lower, upper)
            #计算轮廓
            contours1, _ = cv2.findContours(mask1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            contours2, _ = cv2.findContours(mask2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            mask1 = cv2.cvtColor(mask1, cv2.COLOR_GRAY2BGR)
            mask2 = cv2.cvtColor(mask2, cv2.COLOR_GRAY2BGR)
            frame1_copy = frame1.copy()
            frame2_copy = frame2.copy()
            #计算frame1的最大轮廓
            max_area = 100
            if contours1 is not None:
                for contour1 in contours1:
                    area1 = cv2.contourArea(contour1)
                    if area1 > max_area:
                        max_area = area1
                        x1, y1, w1, h1 = cv2.boundingRect(contour1)
                if max_area > 100:      
                    cv2.rectangle(frame1_copy, (x1, y1), (x1 + w1, y1 + h1), (0, 255, 0), 2)
                    #cv2.rectangle(mask1, (x1, y1), (x1 + w1, y1 + h1), (0, 255, 0), 2)
            #计算frame2的最大轮廓
            max_area = 100
            if contours2 is not None:
                for contour2 in contours2:
                    area2 = cv2.contourArea(contour2)
                    if area2 > max_area:
                        max_area = area2
                        x2, y2, w2, h2 = cv2.boundingRect(contour2)
                if max_area > 100:      
                    cv2.rectangle(frame2_copy, (x2, y2), (x2 + w2, y2 + h2), (0, 255, 0), 2)
                    #cv2.rectangle(mask2, (x2, y2), (x2 + w2, y2 + h2), (0, 255, 0), 2)
            # ------------------------------------双目测距与目标检测----------------------------------------------------------
            imgL = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)  # 将BGR格式转换成灰度图片
            imgR = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

            # cv2.remap 重映射，就是把一幅图像中某位置的像素放置到另一个图片指定位置的过程。
            # 依据MATLAB测量数据重建无畸变图片
            img1_rectified = cv2.remap(imgL, left_map1, left_map2, cv2.INTER_LINEAR)
            img2_rectified = cv2.remap(imgR, right_map1, right_map2, cv2.INTER_LINEAR)

            imageL = cv2.cvtColor(img1_rectified, cv2.COLOR_GRAY2BGR)
            imageR = cv2.cvtColor(img2_rectified, cv2.COLOR_GRAY2BGR)

            # ------------------------------------SGBM算法----------------------------------------------------------
            #   blockSize                   深度图成块，blocksize越低，其深度图就越零碎，0<blockSize<10
            #   img_channels                BGR图像的颜色通道，img_channels=3，不可更改
            #   numDisparities              SGBM感知的范围，越大生成的精度越好，速度越慢，需要被16整除，如numDisparities
            #                               取16、32、48、64等
            #   mode                        sgbm算法选择模式，以速度由快到慢为：STEREO_SGBM_MODE_SGBM_3WAY、
            #                               STEREO_SGBM_MODE_HH4、STEREO_SGBM_MODE_SGBM、STEREO_SGBM_MODE_HH。精度反之
            # -----------------------------------------------------------------------------------------------------
            blockSize = 8
            img_channels = 3
            stereo = cv2.StereoSGBM_create(minDisparity=1,
                                        numDisparities=16,
                                        blockSize=blockSize,
                                        P1=8 * img_channels * blockSize * blockSize,
                                        P2=32 * img_channels * blockSize * blockSize,
                                        disp12MaxDiff=-1,
                                        preFilterCap=1,
                                        uniquenessRatio=10,
                                        speckleWindowSize=100,
                                        speckleRange=100,
                                        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY)
            # 计算视差
            disparity = stereo.compute(img1_rectified, img2_rectified)
            # 归一化函数算法
            disp = cv2.normalize(disparity, disparity, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            # 计算三维坐标数据值
            threeD = cv2.reprojectImageTo3D(disparity, Q, handleMissingValues=True)
            threeD = threeD * 16

            # 格式转变，BGRtoRGB
            frame1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2RGB)
            # 转变成Image格式
            frame1 = Image.fromarray(np.uint8(frame1))
            frame1_shape = np.array(np.shape(frame1)[0:2])
            # 调整图片大小、颜色通道，使其适应YOLO推理的格式
            # frame1 = resize_image(frame1,(640,480))
            frame1 = cv2.cvtColor(np.array(frame1), cv2.COLOR_RGB2BGR)

            # 推理图片
            result = det.predict(frame1)
            # 画框，标出识别的类别、距离、置信度等
            frame1 = visualize(frame1, result)
            #画框，颜色阈值发现的障碍物
            if contours1 is not None and contours2 is not None:
                x_mean = int((x1+x1+w1+x2+x2+w2)/4)
                y_mean = int((y1+y1+h1+y2+y2+h2)/4)
                cv2.circle(frame1_copy,(x_mean,y_mean),radius=10,color=(255,0,0),thickness=-1)   
                cv2.circle(frame2_copy,(x_mean,y_mean),radius=10,color=(255,0,0),thickness=-1)
                distance = math.sqrt(threeD[y_mean][x_mean][0] ** 2 + threeD[y_mean][x_mean][1] ** 2 + threeD[y_mean][x_mean][2] ** 2)
                distance = distance / 1000
                cv2.putText(frame1_copy, "%.2fm"%distance, (x_mean, y_mean), cv2.FONT_HERSHEY_SIMPLEX,1, (0, 255, 0), 2)
                cv2.putText(frame2_copy, "%.2fm"%distance, (x_mean, y_mean), cv2.FONT_HERSHEY_SIMPLEX,1, (0, 255, 0), 2)
                print("obstacles distance:%.2fm"%distance)
            img_show = cv2.hconcat([frame1_copy, frame2_copy])
            cv2.imshow("img_show", img_show)
            # 计算帧率
            fps = (fps + (1. / (time.time() - t1))) / 2
            frame1 = cv2.putText(frame1, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.imshow("frame1", frame1)
            cv2.imshow(WIN_NAME, disp)  # 显示深度图的双目画面
            # 若键盘按下q则退出播放
            if cv2.waitKey(1) & 0xff == ord('q'):
                break

    # 4 释放资源
    det.free()
    cap.release()

    # 5 关闭所有窗口
    cv2.destroyAllWindows()
